#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
SM2Algorithm特定的Side-Channel Attack实现
"""

import numpy as np
from typing import Dict, Any, List
from ..base_algorithm_attack import AlgorithmSpecificAttack
from .sm2_models import SM2LeakageModels


class SM2SpecificAttacks(AlgorithmSpecificAttack):
    """SM2特定的Side-Channel Attack"""
    
    def __init__(self):
        super().__init__("SM2")
        
        # SM2Parameters
        self.key_size = 256  # SM2Use256位Key
        self.block_size = 32  # 256位 = 32字节
        self.round_count = 1   # SM2没Has轮的概念
        
        # 新架构：分离的模型
        self.intermediate_functions = SM2LeakageModels.get_intermediate_functions()
        self.leakage_models = SM2LeakageModels.get_leakage_models()
        self.operation_mapping = SM2LeakageModels.get_operation_to_intermediate_mapping()
        
        # SM2曲线Parameters
        self.curve_p = SM2LeakageModels.SM2_P
        self.curve_a = SM2LeakageModels.SM2_A
        self.curve_b = SM2LeakageModels.SM2_B
        self.curve_n = SM2LeakageModels.SM2_N
        self.base_x = SM2LeakageModels.SM2_GX
        self.base_y = SM2LeakageModels.SM2_GY
        
    def get_algorithm_parameters(self) -> Dict[str, Any]:
        """获取SM2AlgorithmParameters"""
        return {
            'algorithm': 'SM2',
            'type': 'asymmetric',
            'key_size': self.key_size,
            'block_size': self.block_size,
            'round_count': self.round_count,
            'curve_p': self.curve_p,
            'curve_a': self.curve_a,
            'curve_b': self.curve_b,
            'curve_n': self.curve_n,
            'base_point': (self.base_x, self.base_y)
        }
    
    def get_intermediate_functions(self) -> Dict[str, Any]:
        """获取Intermediate valueCalculateFunction"""
        return self.intermediate_functions
    
    def get_leakage_models(self) -> Dict[str, Any]:
        """获取泄漏表示模型"""
        return self.leakage_models
    
    def create_leakage_models(self) -> Dict[str, Any]:
        """CreateSM2特定的泄漏模型（用于基Class兼容）"""
        combined = {}
        combined.update(self.intermediate_functions)
        combined.update(self.leakage_models)
        return combined
    
    def get_target_operations(self) -> List[str]:
        """获取SM2可Attack的目标操作"""
        return SM2LeakageModels.get_target_operations()
    
    def get_operation_to_intermediate_mapping(self) -> Dict[str, str]:
        """获取目标操作ToIntermediate valueFunction的Mapping"""
        return self.operation_mapping
    
    def _calculate_intermediate_values(self, target_data: np.ndarray, hypothesis: int, 
                                     target_operation: str) -> np.ndarray:
        """
        CalculateSM2特定的Intermediate value
        
        Args:
            target_data: 目标Data（Plaintext字节）
            hypothesis: False设值（Key字节的False设）
            target_operation: 目标操作
            
        Returns:
            Intermediate valueArray（预测的汉明权重）
        """
        if target_operation == 'scalar_multiplication':
            # 改进的SM2Scalar乘法泄漏模型（与Power仿True模型Match）
            results = []
            for td in target_data:
                msg_byte = int(td) & 0xFF  # EnsureYes0-255范围
                key_byte = hypothesis & 0xFF  # Key字节False设
                
                # CalculateIntermediate value（与Power仿True一致）
                intermediate = (msg_byte * key_byte) % 256
                
                # Calculate汉明权重（与Power模型一致）
                hw_intermediate = bin(intermediate).count('1')
                hw_key = bin(key_byte).count('1')
                
                # 预测的Power泄漏值（与Power仿True的hw_factor公式对应）
                leakage = hw_intermediate * 0.5 + hw_key * 0.4
                results.append(leakage)
            
            return np.array(results)
        
        elif target_operation == 'point_addition':
            # SM2点加法泄漏模型
            results = []
            for td in target_data:
                x1, y1 = int(td) % 100, hypothesis % 100
                x2, y2 = (int(td) + 1) % 100, (hypothesis + 1) % 100
                try:
                    x3, y3 = SM2LeakageModels.sm2_point_addition_model(x1, y1, x2, y2, 101)
                    results.append(x3)
                except:
                    results.append(0)
            return np.array(results)
        
        elif target_operation == 'sm2_signing':
            # SM2签名泄漏模型
            results = []
            for td in target_data:
                message_hash = int(td) % 1000
                private_key = hypothesis % 100
                nonce = (int(td) + hypothesis) % 100 + 1  # Avoid0
                try:
                    r, s = SM2LeakageModels.sm2_sign_model(
                        message_hash, private_key, nonce, 
                        self.base_x % 1000, self.base_y % 1000, 997, 991)
                    results.append(r)
                except:
                    results.append(0)
            return np.array(results)
        
        elif target_operation == 'sm2_encryption':
            # SM2加密泄漏模型
            results = []
            for td in target_data:
                message = int(td) % 1000
                pub_x, pub_y = hypothesis % 100, (hypothesis + 1) % 100
                nonce = (int(td) + hypothesis) % 100 + 1
                try:
                    c1_x, c1_y, c2 = SM2LeakageModels.sm2_encrypt_model(
                        message, pub_x, pub_y, nonce, 
                        self.base_x % 1000, self.base_y % 1000, 997)
                    results.append(c1_x)
                except:
                    results.append(0)
            return np.array(results)
        
        elif target_operation == 'sm3_hashing':
            # SM3哈希泄漏模型
            results = []
            for td in target_data:
                # 简化的消息块
                message_block = [int(td) % 1000] * 16
                hash_result = SM2LeakageModels.sm3_compression_model(message_block)
                results.append(hash_result[0])  # 取First个字
            return np.array(results)
        
        else:
            # DefaultUseScalar乘法模型
            return self._calculate_intermediate_values(target_data, hypothesis, 'scalar_multiplication')
    
    def _get_selection_function(self, target_data: np.ndarray, hypothesis: int,
                              target_operation: str) -> np.ndarray:
        """
        获取SM2特定的选择Function（用于DPAAttack）
        
        Args:
            target_data: 目标Data
            hypothesis: False设值
            target_operation: 目标操作
            
        Returns:
            选择FunctionResult（0Or1的Array）
        """
        intermediate_values = self._calculate_intermediate_values(
            target_data, hypothesis, target_operation
        )
        
        if target_operation in ['scalar_multiplication', 'point_addition']:
            # Use坐标值的最低位
            return intermediate_values & 1
        
        elif target_operation == 'sm2_signing':
            # Use签名值的汉明权重奇偶性
            hamming_weights = np.array([SM2LeakageModels.hamming_weight(val) 
                                      for val in intermediate_values])
            return hamming_weights & 1
        
        elif target_operation == 'sm2_encryption':
            # Use加密Result的最高位
            return (intermediate_values >> 7) & 1
        
        elif target_operation == 'sm3_hashing':
            # Use哈希值的最低位
            return intermediate_values & 1
        
        else:
            # DefaultUse最低位
            return intermediate_values & 1
    
    def attack_scalar_multiplication(self, **kwargs):
        """
        AttackSM2Scalar乘法操作
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        return self.run_cpa_attack(
            target_operation='scalar_multiplication',
            **kwargs
        )
    
    def attack_sm2_signing(self, **kwargs):
        """
        AttackSM2签名过程
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        return self.run_dpa_attack(
            target_operation='sm2_signing',
            **kwargs
        )
    
    def attack_sm2_encryption(self, **kwargs):
        """
        AttackSM2加密过程
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        return self.run_template_attack(
            target_operation='sm2_encryption',
            **kwargs
        )
    
    def attack_sm3_hashing(self, **kwargs):
        """
        AttackSM3哈希过程
        
        Args:
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        return self.run_spa_attack(
            target_operation='sm3_hashing',
            **kwargs
        )
    
    def private_key_recovery_attack(self, key_bit_length: int = 16, **kwargs) -> Dict[str, Any]:
        """
        SM2私钥恢复Attack
        
        Args:
            key_bit_length: 要恢复的私钥位长度
            **kwargs: AttackParameters
            
        Returns:
            AttackResult
        """
        results = {}
        recovered_key_bits = {}
        
        self.logger.info(f"Starting SM2 private key recovery attack ({key_bit_length} bits)")
        
        # Attack每一位私钥
        for bit_pos in range(key_bit_length):
            self.logger.info(f"Attacking bit {bit_pos}")
            
            try:
                # UseCPAAttackScalar乘法
                result = self.attack_scalar_multiplication(
                    target_byte=bit_pos % 16,  # MappingTo字节位置
                    **kwargs
                )
                
                results[f"bit_{bit_pos}"] = result
                
                if result.success and result.recovered_key_bytes:
                    # 将恢复的字节ConvertFor位
                    recovered_bit = result.recovered_key_bytes[0] & 1
                    recovered_key_bits[bit_pos] = recovered_bit
                    self.logger.info(f"Bit {bit_pos}: {recovered_bit} "
                                   f"(confidence: {result.confidence:.4f})")
                else:
                    self.logger.warning(f"Failed to recover bit {bit_pos}")
            
            except Exception as e:
                self.logger.error(f"Error attacking bit {bit_pos}: {e}")
        
        # 构建恢复的私钥位串
        recovered_key_bits_list = []
        for i in range(key_bit_length):
            if i in recovered_key_bits:
                recovered_key_bits_list.append(recovered_key_bits[i])
            else:
                recovered_key_bits_list.append(0)  # 未恢复的位设For0
        
        # 将位串ConvertForInteger
        recovered_key_int = 0
        for i, bit in enumerate(recovered_key_bits_list):
            recovered_key_int |= (bit << i)
        
        success_rate = len(recovered_key_bits) / key_bit_length
        
        summary = {
            'algorithm': 'SM2',
            'attack_type': 'Private Key Recovery',
            'total_bits': key_bit_length,
            'recovered_bits': len(recovered_key_bits),
            'success_rate': success_rate,
            'recovered_key_bits': recovered_key_bits_list,
            'recovered_key_int': recovered_key_int,
            'recovered_key_hex': f'{recovered_key_int:0{(key_bit_length+3)//4}X}',
            'individual_results': results
        }
        
        self.logger.info(f"SM2 key recovery completed: {len(recovered_key_bits)}/{key_bit_length} bits recovered "
                        f"({success_rate:.1%} success rate)")
        
        return summary
    
    def validate_recovered_key(self, recovered_key_int: int, 
                             test_message: bytes, test_signature: tuple,
                             public_key: tuple) -> bool:
        """
        Verify恢复的私钥YesNo正确
        
        Args:
            recovered_key_int: 恢复的私钥Integer
            test_message: Test消息
            test_signature: Test签名 (r, s)
            public_key: 公钥点 (x, y)
            
        Returns:
            私钥YesNo正确
        """
        try:
            # 简化的SM2Verify
            # 实际实现RequireComplete的SM2VerifyAlgorithm
            from ...crypto.sm2_signature import SM2Signature
            
            # 这里可以CallSM2VerifyFunction
            # 简化ReturnTrue
            return True
            
        except Exception as e:
            self.logger.error(f"Key validation failed: {e}")
            return False
